rm(list=ls())
library(fields); library(maps); library(ncdf4);library(abind)
library(geosphere);library(mapdata)
setwd('/Volumes/Disk2/JJA_ozone_MAM/Harvard_Dataverse')

#----- load functions -----------------------
source("./Function/get_geo.R")
source("./Function/get_met.R")
source('./Function/read_detrend.R')
source('./Function/eof.mca.R')
source('./Function/cov4gappy.R')
source('./Function/Henderson.R')
source('./Function/filled.contour3.R')

#------ read ozone -------------------------
load("./Data/regrid_JJA-yearly-O3_1980-2013.Rdata")
for (i in 1:28){
	for (j in 1:11){
		ap= mov.detrend(yearly_O3[i,j,],len=7)
		yearly_O3[i,j,]=ap
	}
}

#==== calculate the avg ozone in the eastern US =============
ind1=(lon.frame>=-100 & lon.frame<=-65)
ind2=(lat.frame>=32.5 & lat.frame<=47.5)
SE.O3=apply(yearly_O3[ind1,ind2,],c(3),mean,na.rm=TRUE)

US.mask=!is.na(apply(yearly_O3,c(1,2),mean,na.rm=TRUE))
US.area=cal.area(lon.frame,lat.frame)
US.area[!US.mask]=NA
y=cal.area_mean(yearly_O3,US.area,ind1,ind2)
SE.O3=y

####start the plot #######
mai=c(0.1, 0.05, 0.2, 0.05)
mgp=c(1.2, 0.4, 0)
tcl=-0.2
ps=12

start_months=c(12, 3,  6)
end_months=c(2,  5,  8)

close.screen(all.screens = TRUE)
m <- rbind(c(0.2, 0.8, 0.85, 1),
				  c(0.03, 0.35, 0.6-0.05, 1-0.05), c(0.35, 0.68, 0.6-0.05, 1-0.05), c(0.68, 1.0, 0.6-0.05, 1-0.05), 
				  c(0.2, 0.8, 0.45-0.01, 0.6-0.01),
				  c(0.03, 0.35, 0.15-0.01, 0.55-0.01), c(0.35, 0.68, 0.15-0.01, 0.55-0.01), c(0.68, 1.0, 0.15-0.01, 0.55-0.01),
				  c(0.2,0.8,0.06-0.01,0.4-0.01))
split.screen(m)

sub_names=c("a. DJF","b. MAM","c. JJA","d. DJF","e. MAM","f. JJA")

col.gen = colorRampPalette(c("#0563FA", "white","#ff4d4d"), space="rgb")
levels=seq(-0.65, 0.65, length=30)
mid=(levels[1:(length(levels)-1)]+levels[2:(length(levels))])/2
cols=col.gen(length(levels)-1)
if(min(abs(mid))>0){
	ind1=which(mid<=0);ind2=which(mid>=0)
	ind=abind(Reduce(intersect, list(ind1+1, ind2)), Reduce(intersect, list(ind1, ind2-1)))
	ind=sort(unique(ind))
  	 cols[ind]="white"
}	

screen(1)
text(x=0.5,y=0.6,"Correlation of SSTs with East-JJA-O3", cex=1.2,font=1)

for (kk in 1:3){
screen(kk+1)	
met=read_SST('sst.mnmean.nc', start_months[kk], end_months[kk],xlim=c(180,360),ylim=c(0,70))

lon=met$lon
lat=met$lat
month_SST=array(NA,dim(met$data[,,33:66]))
for (i in 1:dim(month_SST)[1]){
	for (j in 1:dim(month_SST)[2]){
		month_SST[i,j,]= mov.detrend(met$data[i,j,][33:66],len=7)
	}
}
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
result=find.cor3(month_SST,SE.O3)
# save("result","lon","lat",file="./Figures/Figure 1/MAM_Ozone_SST.Rdata")

ap=result$cor
p=result$p
ap[ap>=0.65]=0.65;ap[ap<=-0.65]=-0.65
filled.contour3(lon-360,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world", add = T)
cr=0.05
clines=contourLines(lon-360, lat, p,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}

if(kk==1){
	arrows(x0=-100,y0=31.25,x1=-80,y1=31.25, code=0,col=4,lwd=1.5)
	arrows(x0=-80,y0=31.25,x1=-65,y1=49, code=0,col=4,lwd=1.5)
	arrows(x0=-100,y0=49,x1=-65,y1=49, code=0,col=4,lwd=1.5)	
	arrows(x0=-100,y0=31.25,x1=-100,y1=49, code=0,col=4,lwd=1.5)		
	mtext("SST",side=2,font=1,cex=1.05)
}
text(sub_names[kk], x=-158,y=4,font=1)

if(kk==2){
	loc=c(-160,-130,32,50)
	arrows(x0=loc[1],y0=loc[3],x1=loc[2],code=0,col=2,lwd=1.5)
	arrows(x0=loc[1],y0=loc[4],x1=loc[2],code=0,col=2,lwd=1.5)
	arrows(x0=loc[1],y0=loc[3],y1=loc[4],code=0,col=2,lwd=1.5)
	arrows(x0=loc[2],y0=loc[3],y1=loc[4],code=0,col=2,lwd=1.5)

	loc=c(-90,-30,5,22)
	arrows(x0=loc[1],y0=loc[3],x1=loc[2],code=0,col=1,lwd=1.5)
	arrows(x0=loc[1],y0=loc[4],x1=loc[2],code=0,col=1,lwd=1.5)
	arrows(x0=loc[1],y0=loc[3],y1=loc[4],code=0,col=1,lwd=1.5)
	arrows(x0=loc[2],y0=loc[3],y1=loc[4],code=0,col=1,lwd=1.5)
}
}

screen(5)
text(x=0.5,y=0.6,"Correlation of SLPs with East-JJA-O3", cex=1.2,font=1)

####################################
for (kk in 1:3){
screen(kk+5)	
met=read_surface('slp.mon.mean.nc',start_months[kk], end_months[kk],xlim=c(180,360),ylim=c(0,70))
lon=met$lon
lat=met$lat
month_SST=array(NA,dim(met$data[,,33:66]))
for (i in 1:dim(month_SST)[1]){
	for (j in 1:dim(month_SST)[2]){
		month_SST[i,j,]=mov.detrend(met$data[i,j,33:66],len=7)
	}
}

par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
result=find.cor3(month_SST,SE.O3)
# save("result","lon","lat",file="./Figures/Figure 1/MAM_Ozone_SLP.Rdata")
ap=result$cor
p=result$p
ap[ap>=0.65]=0.65;ap[ap<=-0.65]=-0.65
filled.contour3(lon-360,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world", add = T)
cr=0.05
clines=contourLines(lon-360, lat, p,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(sub_names[kk+3], x=-158,y=4.5,font=1)
if(kk==1) mtext("SLP",side=2,font=1,cex=1.05)

if(kk==2){
	loc=c(-160,-130,10,32.5)
	arrows(x0=loc[1],y0=loc[3],x1=loc[2],code=0,col=1,lwd=1.5)
	arrows(x0=loc[1],y0=loc[4],x1=loc[2],code=0,col=1,lwd=1.5)
	arrows(x0=loc[1],y0=loc[3],y1=loc[4],code=0,col=1,lwd=1.5)
	arrows(x0=loc[2],y0=loc[3],y1=loc[4],code=0,col=1,lwd=1.5)

	loc=c(-110,-70,12,40)
	arrows(x0=loc[1],y0=loc[3],x1=loc[2],code=0,col=2,lwd=1.5)
	arrows(x0=loc[1],y0=loc[4],x1=loc[2],code=0,col=2,lwd=1.5)
	arrows(x0=loc[1],y0=loc[3],y1=loc[4],code=0,col=2,lwd=1.5)
	arrows(x0=loc[2],y0=loc[3],y1=loc[4],code=0,col=2,lwd=1.5)
}

}

screen(9)
mai=c(0.1, 0.1, 0.05, 0.2)
par(mai=mai, mgp=mgp, tcl=tcl, ps=12)
image.plot(ap,legend.shrink=0.9,legend.only=TRUE,zlim=c(-0.65,0.65),col=cols,horizontal=TRUE,legend.width=2.5,axis.args = list(mgp = c(0, 0.8, 0),tcl=-0.2,padj=-1))
text("r",x=par("usr")[2]-0.23,y=par("usr")[3]+0.06,srt=0, adj = 0,xpd=TRUE,cex=1.4)

